import numpy as np
from train_config import get_candidates, divide_candidates, get_statistical_function_and_logging_str, get_info
from easy_io import read_pkl_file, write_pkl_file
from itertools import chain
import os


def predict(candidates, probs):
    probs = np.asarray(probs)
    assert probs.shape == (len(candidates), 2)
    for c, p in zip(candidates, probs):
        c.update(prob=float(p[1]))
    return candidates


def main(candidate_file, result_files, outfile):
    candidates = get_candidates(candidate_file)
    prob_lst = []
    for result_file in result_files:
        valid_probs = read_pkl_file(result_file)
        prob_lst.append(valid_probs)
    valid_probs = np.mean(prob_lst, axis=0)
    candidates = predict(candidates, valid_probs)
    write_pkl_file(outfile, candidates)


if __name__ == '__main__':
    # main(
    #     candidate_file='kaggle_beni_candidates_v2.pkl',
    #     result_files=[
    #         '/data_4t/huzq/kaggle_results/cddv5_fold0_kaggle_beni/cddv5_fold0_kaggle_beni.pkl',
    #         '/data_4t/huzq/kaggle_results/cddv5_fold1_kaggle_beni/cddv5_fold1_kaggle_beni.pkl',
    #         '/data_4t/huzq/kaggle_results/cddv5_fold2_kaggle_beni/cddv5_fold2_kaggle_beni.pkl',
    #         '/data_4t/huzq/kaggle_results/cddv5_fold3_kaggle_beni/cddv5_fold3_kaggle_beni.pkl',
    #         '/data_4t/huzq/kaggle_results/cddv5_fold4_kaggle_beni/cddv5_fold4_kaggle_beni.pkl',
    #     ],
    #     outfile='/data_4t/Kaggle/candidates/3dcnn_kaggle_beni_candidates_list_v2.pkl',
    # )
    # main(
    #     candidate_file='kaggle_testset_candidates_v2.pkl',
    #     result_files=[
    #         '/data_4t/huzq/kaggle_results/cddv5_fold0_kaggle_testset/cddv5_fold0_kaggle_testset.pkl',
    #         '/data_4t/huzq/kaggle_results/cddv5_fold1_kaggle_testset/cddv5_fold1_kaggle_testset.pkl',
    #         '/data_4t/huzq/kaggle_results/cddv5_fold2_kaggle_testset/cddv5_fold2_kaggle_testset.pkl',
    #         '/data_4t/huzq/kaggle_results/cddv5_fold3_kaggle_testset/cddv5_fold3_kaggle_testset.pkl',
    #         '/data_4t/huzq/kaggle_results/cddv5_fold4_kaggle_testset/cddv5_fold4_kaggle_testset.pkl',
    #     ],
    #     outfile='/data_4t/Kaggle/candidates/3dcnn_kaggle_testset_candidates_list_v2.pkl',
    # )

    # cddnames = ['spie_candidates_v4.pkl', 'kaggle_testset_candidates_v4.pkl', 'kaggle_beni_candidates_v4.pkl']
    # for cddname in cddnames:
    #     main(
    #         candidate_file='/ssd_1t/huzq/kaggle_data/'+cddname,
    #         result_files=[
    #             '/data_4t/huzq/kaggle_results/cddv9_fold{i}_{cddname}/cddv9_fold{i}_{cddname}.pkl'.format(i=i, cddname=os.path.splitext(cddname)[0]) for i in range(5)
    #         ],
    #         outfile='/data_4t/Kaggle/candidates/3dcnn_'+cddname,
    #     )

    # cddnames = ['lidc_kaggle_candidates_v9_nomerge.pkl']
    # for cddname in cddnames:
    #     main(
    #         candidate_file='/ssd_1t/huzq/kaggle_data/' + cddname,
    #         result_files=[
    #             '/data_4t/huzq/kaggle_results/cddv9_fold{i}_{cddname}/cddv9_fold{i}_{cddname}.pkl'.format(i=i, cddname=
    #             os.path.splitext(cddname)[0]) for i in range(5)
    #             ],
    #         outfile='/data_4t/Kaggle/candidates/3dcnn_' + cddname,
    #     )

    # cddnames = ['spie_candidates_v4.pkl', 'kaggle_testset_candidates_v4.pkl', 'kaggle_beni_candidates_v4.pkl']
    # for cddname in cddnames:
    #     main(
    #         candidate_file='/ssd_1t/huzq/kaggle_data/' + cddname,
    #         result_files=[
    #             '/data_4t/huzq/kaggle_results/cddv9_vgg13shortcutsv2_bs16_fold{i}_{cddname}/cddv9_vgg13shortcutsv2_bs16_fold{i}_{cddname}.pkl'.format(i=i, cddname=
    #             os.path.splitext(cddname)[0]) for i in range(5)
    #             ],
    #         outfile='/data_4t/Kaggle/candidates/3dcnn_cddv9_vgg13shortcutsv2_bs16_' + cddname,
    #     )
    #
    # cddnames = ['spie_candidates_v4.pkl', 'kaggle_testset_candidates_v4.pkl', 'kaggle_beni_candidates_v4.pkl']
    # for cddname in cddnames:
    #     main(
    #         candidate_file='/ssd_1t/huzq/kaggle_data/' + cddname,
    #         result_files=[
    #             '/data_4t/huzq/kaggle_results/cddv9_vgg10_fold{i}_{cddname}/cddv9_vgg10_fold{i}_{cddname}.pkl'.format(
    #                 i=i, cddname=
    #                 os.path.splitext(cddname)[0]) for i in range(5)
    #             ],
    #         outfile='/data_4t/Kaggle/candidates/3dcnn_cddv9_vgg10_' + cddname,
    #     )

    # cddnames = ['lidc_kaggle_candidates_v13_nomerge.pkl', 'kaggle_beni_candidates_v8_nomerge.pkl', 'kaggle_testset_candidates_v8_nomerge.pkl', 'spie_candidates_v8_nomerge.pkl']
    # for cddname in cddnames:
    #     main(
    #         candidate_file='/ssd_1t/huzq/kaggle_data/' + cddname,
    #         result_files=[
    #             '/data_4t/huzq/kaggle_results/cddv9onv13_fold{i}_{cddname}/cddv9onv13_fold{i}_{cddname}.pkl'.format(
    #                 i=i, cddname=
    #                 os.path.splitext(cddname)[0]) for i in range(5)
    #             ],
    #         outfile='/data_4t/Kaggle/candidates/3dcnn_cddv9onv13_' + cddname,
    #     )

    # cddnames = ['kaggle_beni_candidates_v8_nomerge.pkl', 'kaggle_testset_candidates_v8_nomerge.pkl',
    #             'spie_candidates_v8_nomerge.pkl']
    # for cddname in cddnames:
    #     main(
    #         candidate_file='/ssd_1t/huzq/kaggle_data/' + cddname,
    #         result_files=[
    #             '/data_4t/huzq/kaggle_results/cddv15_v2_fold{i}_{cddname}/cddv15_v2_fold{i}_{cddname}.pkl'.format(
    #                 i=i, cddname=
    #                 os.path.splitext(cddname)[0]) for i in range(4)
    #             ],
    #         outfile='/data_4t/Kaggle/candidates/3dcnn_cddv15_v2_' + cddname,
    #     )

    cddnames = ['kaggle_beni_candidates_v10.pkl', 'kaggle_testset_candidates_v10.pkl', 'spie_candidates_v10.pkl']
    for cddname in cddnames:
        main(
            candidate_file='/ssd_1t/huzq/kaggle_data/' + cddname,
            result_files=[
                '/data_4t/huzq/kaggle_results/cddv15_v2_fold{i}_{cddname}/cddv15_v2_fold{i}_{cddname}.pkl'.format(
                    i=i, cddname=
                    os.path.splitext(cddname)[0]) for i in range(4)
                ],
            outfile='/data_4t/Kaggle/candidates/3dcnn_cddv15_v2_' + cddname,
        )
